package org.lanyu.springainovel.common.repository;

import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Repository;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义JDBC聊天记忆仓库实现
 * 使用自定义的spring_ai_chat_memory表结构，支持content_index排序字段
 * 同时保持与Spring AI ChatMemory框架的兼容�? */
@Repository
public class CustomJdbcChatMemoryRepository {

    private final JdbcTemplate jdbcTemplate;
    private final Map<String, List<Message>> conversationCache = new ConcurrentHashMap<>();
    
    @Autowired
    private org.lanyu.springainovel.chat.mapper.ChatMemoryMapper chatMemoryMapper;

    public CustomJdbcChatMemoryRepository(JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
        initializeTable();
    }

    /**
     * 初始化表结构（如果不存在）
     */
    private void initializeTable() {
        // 由于表已在SQL文件中定义，这里不再创建表，只检查表是否存在
        String checkTableSql = "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = 'spring_ai_chat_memory'";
        Integer tableExists = jdbcTemplate.queryForObject(checkTableSql, Integer.class);
        
        if (tableExists == null || tableExists == 0) {
            // 如果表不存在，则创建表（与SQL文件中的定义一致）
            String createTableSql = """
                CREATE TABLE spring_ai_chat_memory (
                  conversation_id varchar(36) NOT NULL,
                  content text NOT NULL,
                  type varchar(10) NOT NULL,
                  content_Index int NULL DEFAULT NULL COMMENT '气泡排序，用于重新生成、删除指定气泡',
                  timestamp timestamp NOT NULL,
                  CONSTRAINT TYPE_CHECK CHECK (type in ('USER','ASSISTANT','SYSTEM','TOOL'))
                ) ENGINE=InnoDB CHARACTER SET=utf8mb4 COLLATE=utf8mb4_german2_ci COMMENT='对话记录表' ROW_FORMAT=Dynamic
                """;
                
            jdbcTemplate.execute(createTableSql);
        }
    }

    /**
     * 获取会话的最近N条消息
     * 
     * @param conversationId 会话ID
     * @param maxMessages 最大消息数量
     * @return 消息列表
     */
    public List<Message> getMessages(String conversationId, int maxMessages) {
        // 先从缓存中获取
        List<Message> cachedMessages = conversationCache.get(conversationId);
        if (cachedMessages != null) {
            return cachedMessages.size() > maxMessages ? 
                cachedMessages.subList(cachedMessages.size() - maxMessages, cachedMessages.size()) : 
                cachedMessages;
        }

        // 从数据库获取，使用MyBatis-Plus
        List<org.lanyu.springainovel.common.entity.SpringAIChatMemory> chatMemories = chatMemoryMapper.getRecentMessages(conversationId, maxMessages);
        
        // 转换为Spring AI消息对象
        List<Message> messages = new ArrayList<>();
        for (org.lanyu.springainovel.common.entity.SpringAIChatMemory memory : chatMemories) {
            Message message = createMessage(memory.getType(), memory.getContent());
            if (message != null) {
                messages.add(message);
            }
        }
        
        // 更新缓存
        conversationCache.put(conversationId, messages);
        
        return messages;
    }
    
    /**
     * 根据类型创建消息对象
     * 
     * @param type 消息类型
     * @param content 消息内容
     * @return Message对象
     */
    private Message createMessage(String type, String content) {
        if (type == null || content == null) {
            return null;
        }
        
        return switch (type.toLowerCase()) {
            case "user" -> new org.springframework.ai.chat.messages.UserMessage(content);
            case "assistant" -> new org.springframework.ai.chat.messages.AssistantMessage(content);
            case "system" -> new org.springframework.ai.chat.messages.SystemMessage(content);
            default -> throw new IllegalArgumentException("Unknown message type: " + type);
        };
    }

    /**
     * 保存消息到数据库
     * 
     * @param conversationId 会话ID
     * @param messages 消息列表
     */
    public void saveMessages(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }

        // 获取当前时间戳
        LocalDateTime timestamp = LocalDateTime.now();
        
        // 转换为SpringAIChatMemory对象并保存
        List<org.lanyu.springainovel.common.entity.SpringAIChatMemory> chatMemories = new ArrayList<>();
        for (Message message : messages) {
            String messageType = message.getMessageType().getValue();
            // 转换为大写，以匹配SQL文件中的约束
            String dbMessageType = messageType.toUpperCase();
            
            org.lanyu.springainovel.common.entity.SpringAIChatMemory memory = 
                new org.lanyu.springainovel.common.entity.SpringAIChatMemory(
                    conversationId, 
                    message.getText(), 
                    dbMessageType, 
                    timestamp
                );
            
            chatMemories.add(memory);
        }
        
        // 批量保存
        for (org.lanyu.springainovel.common.entity.SpringAIChatMemory memory : chatMemories) {
            chatMemoryMapper.insert(memory);
        }
        
        // 更新缓存
        List<Message> cachedMessages = conversationCache.computeIfAbsent(conversationId, k -> new ArrayList<>());
        cachedMessages.addAll(messages);
    }

    /**
     * 清除会话的所有消息
     * @param conversationId 会话ID
     */
    public void clearMessages(String conversationId) {
        if (conversationId == null) {
            return;
        }
        
        // 从数据库删除
        chatMemoryMapper.deleteByConversationId(conversationId);
        
        // 清除缓存
        conversationCache.remove(conversationId);
    }

    /**
     * 删除指定时间戳及之后的所有消息
     * 
     * @param conversationId 会话ID
     * @param fromTimestamp 起始时间戳
     */
    public void deleteMessagesFromTimestamp(String conversationId, LocalDateTime fromTimestamp) {
        if (conversationId == null || fromTimestamp == null) {
            return;
        }
        
        // 从数据库删除
        chatMemoryMapper.deleteFromTimestamp(conversationId, fromTimestamp);
        
        // 更新缓存 - 重新获取剩余消息
        List<org.lanyu.springainovel.common.entity.SpringAIChatMemory> remainingMemories = 
            chatMemoryMapper.getRecentMessages(conversationId, Integer.MAX_VALUE);
        
        List<Message> remainingMessages = new ArrayList<>();
        for (org.lanyu.springainovel.common.entity.SpringAIChatMemory memory : remainingMemories) {
            Message message = createMessage(memory.getType(), memory.getContent());
            if (message != null) {
                remainingMessages.add(message);
            }
        }
        
        if (remainingMessages.isEmpty()) {
            conversationCache.remove(conversationId);
        } else {
            conversationCache.put(conversationId, remainingMessages);
        }
    }



    /**
     * 获取指定会话的ChatMemory实例
     * 
     * @param conversationId 会话ID
     * @param maxMessages 最大消息数�?     * @return ChatMemory实例
     */
    public ChatMemory getChatMemory(String conversationId, int maxMessages) {
        return new CustomJdbcChatMemory(conversationId, maxMessages, this);
    }

    /**
     * 自定义JDBC聊天记忆实现
     */
    private static class CustomJdbcChatMemory implements ChatMemory {
        private final String conversationId;
        private final int maxMessages;
        private final CustomJdbcChatMemoryRepository repository;

        public CustomJdbcChatMemory(String conversationId, int maxMessages, CustomJdbcChatMemoryRepository repository) {
            this.conversationId = conversationId;
            this.maxMessages = maxMessages;
            this.repository = repository;
        }

        @Override
        public void add(String conversationId, Message message) {
            List<Message> messages = new ArrayList<>();
            messages.add(message);
            // 使用实例的conversationId而不是参数中的conversationId
            repository.saveMessages(this.conversationId, messages);
        }

        @Override
        public void add(String conversationId, List<Message> messages) {
            // 使用实例的conversationId而不是参数中的conversationId
            repository.saveMessages(this.conversationId, messages);
        }

        @Override
        public List<Message> get(String conversationId) {
            // 使用实例的conversationId而不是参数中的conversationId
            return repository.getMessages(this.conversationId, maxMessages);
        }

        @Override
        public void clear(String conversationId) {
            // 使用实例的conversationId而不是参数中的conversationId
            repository.clearMessages(this.conversationId);
        }
    }
}
